import os
import glob
import cv2
import numpy as np

def generate_motion_kernel(kernel_size, angle):
    """
    生成运动模糊核（PSF）

    参数：
      kernel_size: 模糊核尺寸，值越大，运动距离越长
      angle: 模糊角度（单位：度），对应物体运动方向
    返回：
      psf: 归一化的运动模糊核
    """
    # 构造一个横向直线核
    kernel = np.zeros((kernel_size, kernel_size))
    kernel[(kernel_size - 1) // 2, :] = np.ones(kernel_size)
    
    # 根据给定角度旋转核
    M = cv2.getRotationMatrix2D((kernel_size / 2, kernel_size / 2), angle, 1)
    kernel_rotated = cv2.warpAffine(kernel, M, (kernel_size, kernel_size))
    
    # 归一化
    psf = kernel_rotated / np.sum(kernel_rotated)
    return psf

def apply_motion_blur(img, kernel_size=15, angle=0):
    """
    对图像应用运动模糊效果

    参数：
      img: 输入图像（BGR格式）
      kernel_size: 模糊核尺寸
      angle: 模糊角度
    返回：
      blurred: 模糊处理后的图像
    """
    psf = generate_motion_kernel(kernel_size, angle)
    blurred = cv2.filter2D(img, -1, psf)
    return blurred

def process_validation_set(validation_folder, output_folder, kernel_size=15, angle=0):
    """
    批量处理验证集中的图像，应用运动模糊

    参数：
      validation_folder: 验证集图像所在文件夹路径
      output_folder: 模糊处理后图像的保存路径
      kernel_size: 模糊核尺寸
      angle: 模糊角度
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # 这里假设验证集图像格式为 jpg，如有需要可以扩展其他格式
    image_paths = glob.glob(os.path.join(validation_folder, '*.jpg'))
    for img_path in image_paths:
        img = cv2.imread(img_path)
        if img is None:
            continue
        blurred_img = apply_motion_blur(img, kernel_size, angle)
        
        # 保存处理后的图像
        filename = os.path.basename(img_path)
        save_path = os.path.join(output_folder, filename)
        cv2.imwrite(save_path, blurred_img)
        print(f'Processed: {filename}')

if __name__ == "__main__":
    # 设置验证集图像所在的文件夹路径
    validation_folder = r"F:\Dataset\sumdata(yes)\val"  # 替换为实际的验证集路径
    # 设置处理后图像的保存文件夹路径
    output_folder = r"F:\Dataset\sumdata(yes)\blur_val"  # 替换为实际保存路径
    
    # 调整 kernel_size 与 angle 参数，根据实际情况模拟不同的运动模糊效果
    process_validation_set(validation_folder, output_folder, kernel_size=15, angle=0)
